package edge

import (
	"context"
	"fmt"

	"github.com/DataDog/KubeHound/pkg/kubehound/graph/adapter"
	"github.com/DataDog/KubeHound/pkg/kubehound/graph/types"
	"github.com/DataDog/KubeHound/pkg/kubehound/models/converter"
	"github.com/DataDog/KubeHound/pkg/kubehound/models/shared"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/cache"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/storedb"
	"github.com/DataDog/KubeHound/pkg/kubehound/store/collections"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

// TokenMountList represent ounts that grant access to the pod service account tokens that reside
// in /var/lib/kubelet/pods/<uid>/volumes/kubernetes.io~projected/<name>/. Paths are normalized by K8s
// to remove the trailing slash.
var TokenMountList = bson.A{
	"/",
	"/var",
	"/var/lib",
	"/var/lib/kubelet",
	"/var/lib/kubelet/pods",
	primitive.Regex{Pattern: "^/var/lib/kubelet/pods/.*"},
}

func init() {
	Register(&ExploitHostTraverse{}, RegisterDefault)
}

type ExploitHostTraverse struct {
	BaseEdge
}

type exploitTraverseTokenGroup struct {
	Parent primitive.ObjectID `bson:"_id" json:"parent"`
	Child  primitive.ObjectID `bson:"child" json:"child"`
}

func (e *ExploitHostTraverse) Label() string {
	return "EXPLOIT_HOST_TRAVERSE"
}

func (e *ExploitHostTraverse) Name() string {
	return "ExploitHostTraverseToken"
}

func (e *ExploitHostTraverse) AttckTechniqueID() AttckTechniqueID {
	return AttckTechniqueUnsecuredCredentials
}

func (e *ExploitHostTraverse) AttckTacticID() AttckTacticID {
	return AttckTacticCredentialAccess
}

func (e *ExploitHostTraverse) Processor(ctx context.Context, oic *converter.ObjectIDConverter, entry any) (any, error) {
	typed, ok := entry.(*exploitTraverseTokenGroup)
	if !ok {
		return nil, fmt.Errorf("invalid type passed to processor: %T", entry)
	}

	return adapter.GremlinEdgeProcessor(ctx, oic, e.Label(), typed.Parent, typed.Child, map[string]any{
		"attckTechniqueID": string(e.AttckTechniqueID()),
		"attckTacticID":    string(e.AttckTacticID()),
	})
}

func (e *ExploitHostTraverse) Stream(ctx context.Context, store storedb.Provider, c cache.CacheReader,
	process types.ProcessEntryCallback, complete types.CompleteQueryCallback) error {

	volumes := adapter.MongoDB(ctx, store).Collection(collections.VolumeName)

	// Link child volumes ONLY where these have interesting properties. Currently this only supports parent
	// directories of the pod token directory to enable TOKEN_STEAL attacks.
	filter := bson.M{
		"type": shared.VolumeTypeHost,
		"source": bson.M{
			"$in": TokenMountList,
		},
		"runtime.runID":        e.runtime.RunID.String(),
		"runtime.cluster.name": e.runtime.Cluster.Name,
	}

	// Find the volume and associated pod namespace and service account.
	pipeline := []bson.M{
		// Look for the volumes that encapsulate the kubernetes node token directory
		{

			"$match": filter,
		},
		{
			"$lookup": bson.M{
				"as":   "childVolumes",
				"from": "volumes",
				"let": bson.M{
					"nid": "$node_id",
				},
				"pipeline": []bson.M{
					{
						"$match": bson.M{
							"$and": bson.A{
								bson.M{"$expr": bson.M{
									"$eq": bson.A{
										"$node_id", "$$nid",
									},
								}},
								bson.M{"type": shared.VolumeTypeProjected},
								bson.M{"projected_id": bson.M{"$ne": nil}},
							},
							"runtime.runID":        e.runtime.RunID.String(),
							"runtime.cluster.name": e.runtime.Cluster.Name,
						},
					},
					{
						"$project": bson.M{
							"_id": 1,
						},
					},
				},
			},
		},
		{
			"$unwind": "$childVolumes",
		},
		{
			"$project": bson.M{
				"_id":   1,
				"child": "$childVolumes._id",
			},
		},
	}

	cur, err := volumes.Aggregate(ctx, pipeline)
	if err != nil {
		return err
	}
	defer cur.Close(ctx)

	return adapter.MongoCursorHandler[exploitTraverseTokenGroup](ctx, cur, process, complete)

}
